from collections import OrderedDict
import math
import torch
import torch.nn.functional as F
import torch.nn as nn
from torch.utils import model_zoo
import copy
import numpy as np
import models.modules
from torchvision import utils
from torch.nn.init import xavier_uniform_, zeros_

import models.senet as senet
import models.modules as modules

# class DispSENet(nn.Module):
    def __init__(self, num_features=2048, block_channel=[256, 512, 1024, 2048]):

        super(DispSENet, self).__init__()

        self.D = modules.D(num_features)
        self.MFF = modules.MFF(block_channel)
        self.R = modules.R(block_channel)

        self.init_weights()
        # original_model = senet.senet154(pretrained='imagenet')
        original_model = senet.se_resnet50(pretrained='imagenet')
        self.E = modules.E_senet(original_model)



    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                xavier_uniform_(m.weight.data)
                if m.bias is not None:
                    zeros_(m.bias)


    def forward(self, x):
        x_block1, x_block2, x_block3, x_block4 = self.E(x)
        x_decoder = self.D(x_block1, x_block2, x_block3, x_block4)
        # x_mff = self.MFF(x_block1, x_block2, x_block3, x_block4,[x_decoder.size(2),x_decoder.size(3)])
        # out = self.R(torch.cat((x_decoder, x_mff), 1))
        out = x_decoder
        return [out]
